package com.abyss.checkpoint;


import org.apache.flink.api.common.restartstrategy.RestartStrategies;
import org.apache.flink.api.common.state.ListState;
import org.apache.flink.api.common.state.ListStateDescriptor;
import org.apache.flink.api.common.time.Time;
import org.apache.flink.api.common.typeinfo.TypeHint;
import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.runtime.state.FunctionInitializationContext;
import org.apache.flink.runtime.state.FunctionSnapshotContext;
import org.apache.flink.runtime.state.filesystem.FsStateBackend;
import org.apache.flink.streaming.api.CheckpointingMode;
import org.apache.flink.streaming.api.checkpoint.CheckpointedFunction;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.environment.CheckpointConfig;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.source.SourceFunction;

import java.util.concurrent.TimeUnit;

/**
 * SavePointDemo
 * 默认情况下，checkpoint不是持久化的，只用于从故障中恢复作业。当程序被取消时，它们会被删除。但是你可以配置checkpoint被周期性持久化到外部，类似于savepoints。这些外部的checkpoints将它们的元数据输出到外部持久化存储并且当作业失败时不会自动清除。这样，如果你的工作失败了，你就会有一个checkpoint来恢复。
 * //如果程序被cancle，保留以前做的checkpoint
 * env.getCheckpointConfig().enableExternalizedCheckpoints(CheckpointConfig.ExternalizedCheckpointCleanup.RETAIN_ON_CANCELLATION);
 * ExternalizedCheckpointCleanup模式配置当你取消作业时外部checkpoint会产生什么行为:
 * ExternalizedCheckpointCleanup.RETAIN_ON_CANCELLATION: 当作业被取消时，保留外部的checkpoint。注意，在此情况下，您必须手动清理checkpoint状态。
 * ExternalizedCheckpointCleanup.DELETE_ON_CANCELLATION: 当作业被取消时，删除外部化的checkpoint。只有当作业失败时，检查点状态才可用。
 */
public class SavePointDemo {
    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(1);

        // CheckPoint config
        env.enableCheckpointing(5000L, CheckpointingMode.EXACTLY_ONCE);
        env.setStateBackend(new FsStateBackend("file:///data/output/checkpoint"));
        env.setRestartStrategy(RestartStrategies.fixedDelayRestart(5, Time.seconds(5L)));
        env.getCheckpointConfig().enableExternalizedCheckpoints(CheckpointConfig.ExternalizedCheckpointCleanup.RETAIN_ON_CANCELLATION);
        env.getCheckpointConfig().setMinPauseBetweenCheckpoints(500L);
        env.getCheckpointConfig().setCheckpointTimeout(10000L);
        env.getCheckpointConfig().setMaxConcurrentCheckpoints(1);

        DataStreamSource<Integer> source = env.addSource(new MySourceForGenerateNumberAccumulate());

        source.print();
        env.execute();
    }

    public static class MySourceForGenerateNumberAccumulate implements SourceFunction<Integer>, CheckpointedFunction {
        private boolean isRun = true;
        private int counter = 0;
        private ListState<Integer> state;
        @Override
        public void run(SourceContext<Integer> ctx) throws Exception {
            while (this.isRun) {
                counter++;
                ctx.collect(counter);
                TimeUnit.SECONDS.sleep(1L);
            }
        }

        @Override
        public void cancel() { this.isRun = false; }

        @Override
        public void snapshotState(FunctionSnapshotContext context) throws Exception {
            this.state.clear();

            this.state.add(counter);
        }

        @Override
        public void initializeState(FunctionInitializationContext context) throws Exception {
            state = context
                    .getOperatorStateStore()
                    .getListState(new ListStateDescriptor<Integer>("counter", TypeInformation.of(new TypeHint<Integer>() {})));
            for (int i: state.get()) {
                this.counter = i;
            }
        }
    }
}